#ifndef CAFFE_INSIGHT_H
#define CAFFE_INSIGHT_H

#include <string>
#include <opencv2/opencv.hpp>

#include <caffe/blob.hpp>
#include <caffe/common.hpp>
#include <caffe/net.hpp>
#include <caffe/proto/caffe.pb.h>
#include <caffe/util/io.hpp>

class CaffeInsight {
public:
  CaffeInsight(const std::string& modelFile, const std::string& weightFile, const std::string& mean_file = "", const std::string& mean_value = "0,0,0", int batchsize = 2, int platform = 0);
  ~CaffeInsight();
  void LoadMeanMat(const std::string& mean_file, const std::string& mean_value);
  void Predict(const cv::Mat* images, float* result, int batchsize, const char* output_layer);
  void SetGPUMode() { caffe::Caffe::set_mode(caffe::Caffe::GPU); }

private:
  void WrapInputBlob(std::vector<std::vector<cv::Mat>>& images, int batchsize);
  void Preprocess(const cv::Mat* images, int batchsize, std::vector<std::vector<cv::Mat>>& inputBlob);
  void Preprocess(const cv::Mat& img, std::vector<cv::Mat>* input_channels);

  std::shared_ptr<caffe::Net<float> > net;
  caffe::Blob<float>* inputBlob;
  int ichannels;
  cv::Size isize;
  std::string modelFile;
  std::string weightFile;
  int batchsize = 2;
  int platform = 0;
  cv::Mat meanMat;
};


#endif  //CAFFE_INSIGHT_H
